from docx import Document
from docx.shared import Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
import re
from datetime import datetime
import pandas as pd
from typing import Dict, Any, List, Tuple, Optional
import gc

class TemplateEngine:
    """简化版Word模板引擎"""
    
    def __init__(self):
        """初始化模板引擎"""
        self.doc = None
        self.df = None
        self.template_path = None
        
        # 基本变量匹配：{{A1}} 或 {{A1|sum}}
        self.var_pattern = re.compile(r'\{\{([^}|]+)(?:\|([^}]+))?\}\}')
        
        # 支持的函数
        self.functions = {
            'sum': lambda x: float(x.sum()),
            'avg': lambda x: float(x.mean()),
            'max': lambda x: float(x.max()),
            'min': lambda x: float(x.min()),
            'count': len
        }
        
        # 特殊变量
        self.special_vars = {
            'date': lambda: datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        }
    
    def _excel_ref_to_index(self, ref: str) -> Tuple[int, int]:
        """将Excel引用转换为索引"""
        try:
            col = 0
            for c in ref.rstrip('0123456789'):
                col = col * 26 + (ord(c.upper()) - ord('A'))
            row = int(ref[len(ref.rstrip('0123456789')):]) - 1
            return row, col
        except Exception as e:
            raise ValueError(f"无效的Excel引用 {ref}: {str(e)}")
    
    def _get_value(self, ref: str, func: Optional[str] = None) -> Any:
        """获取单元格或区域的值"""
        try:
            # 处理特殊变量
            if ref in self.special_vars:
                return self.special_vars[ref]()
            
            # 处理Excel引用
            if ':' in ref:  # 区域引用
                start_ref, end_ref = ref.split(':')
                start_row, start_col = self._excel_ref_to_index(start_ref)
                end_row, end_col = self._excel_ref_to_index(end_ref)
                
                values = self.df.iloc[start_row:end_row+1, start_col:end_col+1].values
                if func and func in self.functions:
                    return self.functions[func](values)
                return values
            else:  # 单元格引用
                row, col = self._excel_ref_to_index(ref)
                return self.df.iloc[row, col]
        except Exception as e:
            return f"引用错误: {str(e)}"
    
    def _process_paragraph(self, paragraph) -> None:
        """处理段落中的变量"""
        if not paragraph or '{{' not in paragraph.text:
            return
            
        text = paragraph.text
        
        def replace_var(match):
            ref = match.group(1).strip()
            func = match.group(2).strip() if match.group(2) else None
            value = self._get_value(ref, func)
            return str(value)
        
        new_text = self.var_pattern.sub(replace_var, text)
        if new_text != text:
            paragraph.text = new_text
            paragraph.alignment = WD_ALIGN_PARAGRAPH.LEFT
            for run in paragraph.runs:
                run.font.size = Pt(10.5)
    
    def process_template(self, df: pd.DataFrame) -> Tuple[bool, str]:
        """处理模板"""
        try:
            if not self.doc:
                return False, "未加载模板"
            
            self.df = df
            
            # 处理文档内容
            for paragraph in self.doc.paragraphs:
                try:
                    self._process_paragraph(paragraph)
                except Exception as e:
                    return False, f"处理段落失败: {str(e)}"
            
            # 处理表格内容
            for table in self.doc.tables:
                for row in table.rows:
                    for cell in row.cells:
                        for paragraph in cell.paragraphs:
                            try:
                                self._process_paragraph(paragraph)
                            except Exception as e:
                                return False, f"处理表格单元格失败: {str(e)}"
            
            return True, "处理成功"
        except Exception as e:
            return False, f"处理模板失败: {str(e)}"
        finally:
            self.df = None
            gc.collect()
    
    def load_template(self, template_path: str) -> Tuple[bool, List[str]]:
        """加载Word模板"""
        try:
            self.cleanup()
            self.template_path = template_path
            self.doc = Document(template_path)
            variables = self.get_template_variables()
            return True, variables
        except Exception as e:
            self.cleanup()
            return False, [f"加载模板失败: {str(e)}"]
    
    def get_template_variables(self) -> List[str]:
        """获取模板中的所有变量"""
        if not self.doc:
            return []
        
        variables = set()
        # 处理文档段落
        for paragraph in self.doc.paragraphs:
            for match in self.var_pattern.finditer(paragraph.text):
                ref = match.group(1).strip()
                func = match.group(2).strip() if match.group(2) else None
                variables.add(f"{ref}|{func}" if func else ref)
        
        # 处理表格
        for table in self.doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    for paragraph in cell.paragraphs:
                        for match in self.var_pattern.finditer(paragraph.text):
                            ref = match.group(1).strip()
                            func = match.group(2).strip() if match.group(2) else None
                            variables.add(f"{ref}|{func}" if func else ref)
        
        return sorted(list(variables))
    
    def save_document(self, output_path: str) -> Tuple[bool, str]:
        """保存文档"""
        if not self.doc:
            return False, "未加载模板"
        
        try:
            self.doc.save(output_path)
            return True, "保存成功"
        except Exception as e:
            return False, f"保存失败: {str(e)}"
        finally:
            self.cleanup()
    
    def cleanup(self) -> None:
        """清理资源但保留路径"""
        try:
            if hasattr(self, 'doc') and self.doc:
                self.doc = None
            if hasattr(self, 'df') and self.df is not None:
                self.df = None
            
            # 不清理template_path，保留路径信息
            
            # 强制进行垃圾回收
            gc.collect()
        except Exception as e:
            print(f"清理资源时出错: {str(e)}") 